#!/usr/bin/env python3
# -*- coding: utf-8 -*-
from reportlab.lib.units import mm
from reportlab.pdfgen import canvas
import os

from PyPDF2 import PdfFileWriter, PdfFileReader

helpDoc = '''
Add Page Number to PDF file with Python 
Python 给 PDF 添加 页码
usage:
require:
    pip install reportlab pypdf2
    Support both Python2/3, But more recommend Python3

tips:
    * only support A4 size PDF
'''


class PdfAddNumber:
    def __init__(self):
        pass

    def start(self):
        dir_list = os.listdir()
        print(dir_list)
        pdf_name_list = []
        for file in dir_list:
            if file.endswith(".pdf"):
                pdf_name_list.append(file)
        print(pdf_name_list)
        new_file_dir = os.getcwd()
        print(new_file_dir)
        new_file_dir = new_file_dir + '/Output/'
        if not os.path.isdir(new_file_dir):
            os.mkdir(new_file_dir)
        if not pdf_name_list:
            print("文件夹中没有找到pdf文件")
            return 0
        for pdf_name in pdf_name_list:
            print("开始添加%s" % pdf_name)
            self.startWritePdf(pdf_name, new_file_dir)
        print("请按任意键退出")

    def createPagePdf(self, num, tmp):
        c = canvas.Canvas(tmp)
        for i in range(1, num + 1):
            c.drawString((210 // 2) * mm, 4 * mm, str(i))
            c.showPage()
        c.save()

    def startWritePdf(self, file_path, new_file_dir):
        tmp = "__tmp.pdf"
        batch = 0
        output = PdfFileWriter()
        with open(file_path, 'rb') as f:
            pdf_reader = PdfFileReader(f, strict=False)
            pdf_number = pdf_reader.getNumPages()
            if batch == 0:
                batch = -pdf_number
            self.createPagePdf(pdf_number, tmp)
            with open(tmp, 'rb') as f_tmp:
                pdf_number_canvas = PdfFileReader(f_tmp)
                for p in range(pdf_number):
                    if not p % batch and p:
                        new_path = file_path.replace(file_path, new_file_dir + os.path.basename(file_path))
                        with open(new_path, 'wb') as fw:
                            output.write(fw)
                        output = PdfFileWriter()
                    print('page: %d of %d' % (p, pdf_number))
                    page = pdf_reader.getPage(p)
                    number_layer = pdf_number_canvas.getPage(p)

                    page.mergePage(number_layer)
                    output.addPage(page)
                if output.getNumPages():
                    new_path = file_path.replace(file_path, new_file_dir + os.path.basename(file_path))
                    with open(new_path, 'wb') as fw:
                        output.write(fw)
            os.remove(tmp)
            print("添加结束")


if __name__ == "__main__":
    pdf_add_number = PdfAddNumber()
    pdf_add_number.start()
    # os.system("pause")

